﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Digraphs;
using Algoritmia.DataStructures.Maps;
using Algoritmia.Utils;

namespace Algoritmia.Problems.AllShortestPaths
{

    public class FloydWarshallAllShortestPathsFinder<T> : IAllShortestPathsFinder<T,double>
    {
        private readonly Func<IEnumerable<T>, IMap<Tuple<T,T>, double>> _createDistancesMap;
        private readonly EqualityComparison<T> _equalityComparison;
        private readonly Func<IEnumerable<T>, IMap<Tuple<T,T>, T>> _createBackpointersMap;
        private IMap<Tuple<T, T>, T> _backpointers;

        public FloydWarshallAllShortestPathsFinder(Func<IEnumerable<T>,IMap<Tuple<T,T>,double>> createDistancesMap = null,
                                                   Func<IEnumerable<T>,IMap<Tuple<T,T>,T>> createBackpointersMap = null,
                                                   EqualityComparison<T> equalityComparison = null)
        {
            _createDistancesMap = createDistancesMap ?? (it => AlgoritmiaFactory.Maps<Tuple<T,T>, double>.Default);
            _createBackpointersMap = createBackpointersMap ?? (it => AlgoritmiaFactory.Maps<Tuple<T,T>, T>.Default);
            _equalityComparison = AlgoritmiaFactory.GetDefaultEqualityComparison(equalityComparison);
        }

        #region Implementation of IAllShortestPathsFinder<T,double>

        public IEnumerable<KeyValuePair<Tuple<T, T>, double>> Distances(IDigraph<T> g, WeightingFunction<T, double> d)
        {
            var distances = _createDistancesMap(g.V);
            foreach(var v in g.V)
            {
                foreach(var w in g.V)
                {
                    var tple = new Tuple<T, T>(v, w);
                    if ( _equalityComparison(v,w) )
                    {
                        distances[tple] = 0.0;
                    }
                    else if ( g.Successors(v).Contains(w))
                    {
                        distances[tple] = d(v,w);
                    }
                    else
                    {
                        distances[tple] = double.PositiveInfinity;
                    }
                }
            }

            foreach(var vk in g.V)
            {
                foreach(var u in g.V)
                {
                    foreach(var w in g.V)
                    {
                        var tple = new Tuple<T, T>(u, w);
                        distances[tple] = Math.Min(distances[tple],
                                                   distances[new Tuple<T, T>(u, vk)] + distances[new Tuple<T, T>(vk, w)]);
                    }
                }
            }
            return distances;
        }

        public IMap<Tuple<T,T>,T> Backpointers(IDigraph<T> g, WeightingFunction<T, double> d)
        {
            var distances = _createDistancesMap(g.V);
            var backpointers = _createBackpointersMap(g.V);
            foreach (var v in g.V)
            {
                foreach (var w in g.V)
                {
                    var tple = new Tuple<T, T>(v, w);
                    if (_equalityComparison(v, w))
                    {
                        distances[tple] = 0.0;
                    }
                    else if (g.Successors(v).Contains(w))
                    {
                        distances[tple] = d(v, w);
                    }
                    else
                    {
                        distances[tple] = double.PositiveInfinity;
                    }
                    backpointers[tple] = tple.Item1;
                }
            }

            foreach (var vk in g.V)
            {
                foreach (var u in g.V)
                {
                    foreach (var w in g.V)
                    {
                        var tple = new Tuple<T, T>(u, w);
                        var candidate = distances[new Tuple<T, T>(u, vk)] + distances[new Tuple<T, T>(vk, w)];
                        if (  candidate < distances[tple] )
                        {
                            distances[tple] = candidate;
                            backpointers[tple] = vk;
                        }
                    }
                }
            }
            return backpointers;
        }

        public IEnumerable<IEnumerable<T>> ShortestPaths(IDigraph<T> g,  WeightingFunction<T, double> d = null, IMap<Tuple<T,T>,T> backpointers = null)
        {
            if ( d == null && backpointers == null )
            {
                throw new ArgumentException("d or backpointers must be not null");
            }
            _backpointers = backpointers ?? Backpointers(g, d);
            foreach(var edge in _backpointers.Keys())
            {
                var path = new[] {edge.Item1}.Concat(Backtrace(edge.Item1, edge.Item2)).Concat(new[] {edge.Item2}).ToList();
                if ( (from i in Enumerable.Range(0,path.Count-1)
                      select new Tuple<T,T>(path[i],path[i+1])).All(g.E.Contains))
                {
                    yield return path;
                }
            }
        }

        #endregion

        private IEnumerable<T> Backtrace(T u, T w)
        {
            T vk  = _backpointers[new Tuple<T, T>(u, w)];

            if (_equalityComparison(vk,u) )
            {
                return Enumerable.Empty<T>();
            }
            return Backtrace(u, vk).Concat(new[] {vk}).Concat(Backtrace(vk, w));
        }
    }
}
